// Copyright lowRISC contributors (OpenTitan project).
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#ifndef OPENTITAN_SW_DEVICE_TESTS_PENETRATIONTESTS_FIRMWARE_LIB_PENTEST_LIB_H_
#define OPENTITAN_SW_DEVICE_TESTS_PENETRATIONTESTS_FIRMWARE_LIB_PENTEST_LIB_H_

#include "sw/device/lib/dif/dif_base.h"
#include "sw/device/lib/dif/dif_uart.h"
#include "sw/device/lib/ujson/ujson.h"
#include "sw/device/tests/penetrationtests/json/pentest_lib_commands.h"

#include "alert_handler_regs.h"  // Generated.
#include "sensor_ctrl_regs.h"    // Generated.

static const char PENTEST_VERSION[16] = "PENTEST: v1.0.3";

// NOP macros used in different FI tests.
#define NOP1 "nop\n"
#define NOP10 NOP1 NOP1 NOP1 NOP1 NOP1 NOP1 NOP1 NOP1 NOP1 NOP1
#define NOP30 NOP10 NOP10 NOP10
#define NOP100 NOP10 NOP10 NOP10 NOP10 NOP10 NOP10 NOP10 NOP10 NOP10 NOP10
#define NOP1000 \
  NOP100 NOP100 NOP100 NOP100 NOP100 NOP100 NOP100 NOP100 NOP100 NOP100

// Optimized (in terms of number of cycles) trigger high function.
// Consists of:
// - 10 NOPS to make sure that the previous code execution has finished
// - Disassembled `dif_gpio_write(&gpio, kTriggerSwBitIndex, true)`
// - 10 NOPS to give the trigger time to raise
// Use with caution as the trigger pin index could change.
#define PENTEST_ASM_TRIGGER_HIGH     \
  asm volatile(NOP10);               \
  asm volatile("lui a0, 0x1000;");   \
  asm volatile("addi a0, a0, 256;"); \
  asm volatile("lui a1, 0x40040;");  \
  asm volatile("sw a0, 24(a1);");    \
  asm volatile(NOP10);

// Optimized (in terms of number of cycles) trigger low function.
// Consists of:
// - 10 NOPS to make sure that the previous code execution has finished
// - Disassembled `dif_gpio_write(&gpio, kTriggerSwBitIndex, false)`
// - 10 NOPS to give the trigger time to fall
// Use with caution as the trigger pin index could change.
#define PENTEST_ASM_TRIGGER_LOW     \
  asm volatile(NOP10);              \
  asm volatile("lui a0, 0x1000;");  \
  asm volatile("lui a1, 0x40040;"); \
  asm volatile("sw a0, 24(a1);");   \
  asm volatile(NOP10);

typedef struct pentest_registered_alerts {
  uint32_t alerts[3];
} pentest_registered_alerts_t;

typedef struct pentest_sensor_alerts {
  uint32_t alerts[2];
} pentest_sensor_alerts_t;

/**
 * Trigger sources.
 *
 * The trigger signal for a peripheral is based on its clkmgr_aon_idle signal.
 * These values must match the values in chiplevel.sv.tpl.
 */
typedef enum pentest_trigger_source {
  /**
   * Use AES for capture trigger.
   *
   * The trigger signal will go high 320 cycles after `dif_aes_trigger()` is
   * called and remain high until the operation is complete.
   */
  kPentestTriggerSourceAes = 0,
  /**
   * Use HMAC for capture trigger.
   */
  kPentestTriggerSourceHmac = 1,
  /**
   * Use KMAC for capture trigger.
   */
  kPentestTriggerSourceKmac = 2,
  /**
   * Use OTBN for capture trigger.
   */
  kPentestTriggerSourceOtbn = 3,
} pentest_trigger_source_t;

/**
 * Trigger type.
 */
typedef enum pentest_trigger_type {
  /**
   * Use the precise hardware capture trigger gateable by software. If selected,
   * the actual capture trigger is generated based on the clkmgr_aon_idle signal
   * of the peripheral corresponding to selected trigger source.
   *
   * Note that this is available on FPGA only.
   */
  kPentestTriggerTypeHwGated = 0,
  /**
   * Use the fully software controlled capture trigger. If selected, the
   * configured trigger source is not relevant.
   */
  kPentestTriggerTypeSw = 1,
} pentest_trigger_type_t;

/**
 * Trigger mask.
 */
typedef enum pentest_mask_type {
  kPentestTrigger1 = 0x1,
  kPentestTrigger2 = 0x2,
  kPentestTrigger3 = 0x4,
  kPentestTrigger4 = 0x8,
  kPentestTrigger5 = 0x10,
  kPentestTrigger6 = 0x20,
  kPentestTrigger7 = 0x40,
  kPentestTrigger8 = 0x80,
} pentest_trigger_mask_t;

/**
 * Peripherals.
 *
 * Constants below are bitmasks that can be used to specify a set of
 * peripherals.
 *
 * See also: `pentest_peripherals_t`.
 */
typedef enum pentest_peripheral {
  /**
   * EDN.
   */
  kPentestPeripheralEdn = 1 << 0,
  /**
   * CSRNG.
   */
  kPentestPeripheralCsrng = 1 << 1,
  /**
   * Entropy source.
   */
  kPentestPeripheralEntropy = 1 << 2,
  /**
   * AES.
   */
  kPentestPeripheralAes = 1 << 3,
  /**
   * HMAC.
   */
  kPentestPeripheralHmac = 1 << 4,
  /**
   * KMAC.
   */
  kPentestPeripheralKmac = 1 << 5,
  /**
   * OTBN.
   */
  kPentestPeripheralOtbn = 1 << 6,
  /**
   * Peripherals using the IO_DIV4_PERI clock (UART, GPIO, I2C, SPI Dev, ...)
   */
  kPentestPeripheralIoDiv4 = 1 << 7,
  /**
   * Peripherals using the IO_DIV2_PERI clock (SPI Host 1)
   */
  kPentestPeripheralIoDiv2 = 1 << 8,
  /**
   * USB.
   */
  kPentestPeripheralUsb = 1 << 9,
  /**
   * Peripherals using the IO_PERI clock (SPI Host 0)
   */
  kPentestPeripheralIo = 1 << 10,
} pentest_peripheral_t;

/**
 * A set of peripherals.
 *
 * This type is used for specifying which peripherals should be enabled when
 * calling `pentest_init()`. Remaining peripherals will be disabled to reduce
 * noise during SCA.
 *
 * See also: `pentest_peripheral_t`, `pentest_init()`.
 */
typedef uint32_t pentest_peripherals_t;

/**
 * A set of contexts for SW LFSR
 *
 * This type is used for specifying the context of the LFSR, with each context
 * corresponding to a different state variable
 */
typedef enum pentest_lfsr_context {
  /**
   * PRNG used for initial data sharing
   */
  kPentestLfsrMasking = 0,
  /**
   * PRNG used for determining the order of fixed and random measurements
   */
  kPentestLfsrOrder = 1,

} pentest_lfsr_context_t;

/**
 * Configures the entropy complex for OTBN tests.
 *
 * Similar to entropy_testutils_auto_mode_init(), this function inits the
 * entropy complex. However, in comparison to the function available in the
 * testutils, this function maximizes the reseed intervall to 0xffffffff.
 * This is necessary to guarantee a fixed trigger window for OTBN tests.
 *
 * @return OK or error.
 */
status_t pentest_configure_entropy_source_max_reseed_interval(void);

/**
 * Returns the registered alerts in the sensor controller.
 *
 * This function reads out the recov_alert and fatal_alert registers of the
 * sensor controller. This controller collects the alerts of the AST.
 *
 * @return A struct containing which of the alerts were triggered during the
 * test.
 */
pentest_sensor_alerts_t pentest_get_sensor_alerts(void);

/**
 * Clears the recoverable AST alerts in the sensor controller.
 */
void pentest_clear_sensor_recov_alerts(void);

/**
 * Returns the registered alerts.
 *
 * If a fault injection triggered an alert, this function returns the alert ID
 * back to the host. Afterwards, the alert is cleared.
 *
 * @return A struct containing which of the alerts were triggered during the
 * test.
 */
pentest_registered_alerts_t pentest_get_triggered_alerts(void);

/**
 * Configures the alert handler for FI.
 *
 * Set how each alert is handled and how the escalation will take place.
 * The alert classes are set as 0 = A, 1 = B, 2 = C, and 3 = D. It is currently
 * not possible to set an alert to the X (None) class. Escalation signals are
 * mapped as 0 meaning Non-Maskable Interrupt (NMI), 1 meaning wipe secrets, 2
 * meaning scrap state (turn off HW blocks), 3 meaning reset. When stting the
 * signal to 4294967295, that signal is disabled. In this function it is not
 * possible to set a single phase to multiple escalation signals. Accumulation
 * threshold means after how many alerts the escalation starts. Recall that a
 * recoverable alert just brings a single alert, and that a fatal alert
 * increases this threshold with each cycle. Duration cycles indicate how many
 * clock cycles each phase lasts.
 *
 * Note that this function only works if the alert handler is left unlocked.
 *
 * @param input_alert_classes The class for each of the 65 alerts.
 * @param input_alert_en Enabling each alert.
 * @param input_class_en Enable each class.
 * @param accumulation_thresholds The threshold of alerts to start the first
 * escalation phase for each class.
 * @param signals How each phase is mapped to escalation signals.
 * @param duration_cycles How many cycles each phase takes.
 * @param ping_timeout How many cycles a timeout happens for pings between
 * blocks.
 */
void pentest_configure_alert_handler(
    uint32_t input_alert_classes[ALERT_HANDLER_PARAM_N_ALERTS],
    bool input_alert_en[ALERT_HANDLER_PARAM_N_ALERTS],
    bool input_class_en[ALERT_HANDLER_PARAM_N_CLASSES],
    uint16_t accumulation_thresholds[ALERT_HANDLER_PARAM_N_CLASSES],
    uint32_t signals[ALERT_HANDLER_PARAM_N_CLASSES],
    uint32_t duration_cycles[ALERT_HANDLER_PARAM_N_CLASSES],
    uint32_t ping_timeout);

/**
 * Reads out the alert handler.
 *
 * The alert classes map 0 to class A, 1 to class B, 2 to class C, and 3 to
 * class D. Please check
 * opentitan/bazel-opentitan/hw/top_earlgrey/sw/autogen/top_earlgrey.c
 * top_earlgrey_alert_peripheral_t top_earlgrey_alert_for_peripheral for what
 * each alert is.
 *
 * @param alert_classes The mapping of each of the alerts to a class.
 * @param enabled_alerts Checks which of the alerts are enabled.
 * @param enabled_classes Checks whether the classes are enabled.
 * @param accumulation_thresholds The accumulation threshold for each class.
 * @param duration_cycles The duration of each escalation phase.
 * @param escalation_signals_en Whether each escalation signal is enabled.
 * @param escalation_signals_map How each escalation signal is mapped to the
 * escalation phases.
 * @return OK.
 */
status_t pentest_print_config_alert_handler(
    uint32_t alert_classes[ALERT_HANDLER_PARAM_N_ALERTS],
    bool enabled_alerts[ALERT_HANDLER_PARAM_N_ALERTS],
    bool enabled_classes[ALERT_HANDLER_PARAM_N_CLASSES],
    uint32_t accumulation_thresholds[ALERT_HANDLER_PARAM_N_CLASSES],
    uint32_t duration_cycles[ALERT_HANDLER_PARAM_N_PHASES],
    bool escalation_signals_en[ALERT_HANDLER_PARAM_N_PHASES],
    uint32_t escalation_signals_map[ALERT_HANDLER_PARAM_N_PHASES]);

/**
 * Return alert config back to host.
 *
 * This function reads out the fields of the alert handler and returns them
 * back to the host over uJSON.
 *
 * @param uj An initialized uJSON context.
 * @return OK.
 */
status_t pentest_send_alert_config(ujson_t *uj);

/**
 * Reads out the sensor control.
 *
 * sensor_ctrl_fatal returns true if the sensor is fatal and false if it is
 * recoverable.
 *
 * @param sensor_ctrl_en Checks the enablement of the sensors.
 * @param sensor_ctrl_fatal Checks the fatality of the sensors.
 * @return OK.
 */
status_t pentest_print_config_sensor_ctrl(
    bool sensor_ctrl_en[SENSOR_CTRL_PARAM_NUM_ALERT_EVENTS],
    bool sensor_ctrl_fatal[SENSOR_CTRL_PARAM_NUM_ALERT_EVENTS]);

/**
 * Return sensor config back to host.
 *
 * This function reads out the fields of the sensor ctrl and returns them
 * back to the host over uJSON.
 *
 * @param uj An initialized uJSON context.
 * @return OK.
 */
status_t pentest_send_sensor_config(ujson_t *uj);

/**
 * Reads the device ID from the LC.
 *
 * Can be used to categorize different SCA and FI runs.
 *
 * @param device_id A buffer available to store the device id.
 * @return OK or error.
 */
status_t pentest_read_device_id(uint32_t device_id[]);

/**
 * Reads the rom digest from the rom controller.
 *
 * Can be used to version the ROM.
 *
 * @param rom_digest A buffer available to store the rom digest.
 * @return OK or error.
 */
status_t pentest_read_rom_digest(uint32_t rom_digest[]);

/**
 * Configures CPU for SCA and FI penetration tests.
 *
 * If disable_icache is True, this functions disables the instruction cache. If
 * disable_icache is False, the instruction cache config remains at
 * enabled/disabled based on the ROM configuration. This implies that the iCache
 * could be disabled even though disable_icache is False.
 *
 * If disable_dummy_instr is True, this functions disables the dummy
 * instructions. If disable_dummy_instr is False, the dummy instruction config
 * remains at the original config that was set in the ROM. This implies that
 * dummy instructions could disabled even though disable_dummy_instr is False.
 *
 * dummy_instr_count sets the number of dummy instructions per real instruction.
 * This works via the following settings. Set to 0 for dummy instruction every 0
 * - 4 real instructions. Set to 1 for dummy instruction every 0 - 8 real
 * instructions. Set to 3 for dummy instruction every 0 - 16 real instructions.
 * Set to 7 for dummy instruction every 0 - 32 real instructions.
 *
 * For enable_jittery_clock, the function first checks if the write protection
 * for this register is enabled. Then, by setting enable_jittery_clock True or
 * False, the jittery clock is enabled or disabled. When regwen is active, no
 * change happens.
 *
 * For enable_sram_readback, the function first checks if the write protection
 * for this register is enabled. Then, by setting enable_sram_readback True or
 * False, the SRAM readback feature for the main and ret SRAM iss enabled or
 * disabled. When regwen is active, no change happens.
 *
 *
 * @param disable_icache Disable the instruction cache.
 * @param disable_dummy_instr Disable the dummy instructions.
 * @param dummy_instr_count Provide how often to add dummy instructions.
 * @param enable_jittery_clock: Enable/disable the jittery clock.
 * @param enable_sram_readback: Enable/disable the SRAM readback feature.
 * @param clock_jitter_locked: Jittery clock register is locked.
 * @param clock_jitter_en: Jittery clock feature is enabled.
 * @param sram_main_readback_locked: SRAM main readback register is locked.
 * @param sram_ret_readback_locked: SRAM ret readback register is locked.
 * @param sram_main_readback_en: SRAM main readback feature is enabled.
 * @param sram_ret_readback_en: SRAM ret readback feature is enabled.
 * @return OK or error.
 */
status_t pentest_configure_cpu(
    bool enable_icache, bool *icache_en, bool enable_dummy_instr,
    bool *dummy_instr_en, uint32_t dummy_instr_count, bool enable_jittery_clock,
    bool enable_sram_readback, bool *clock_jitter_locked, bool *clock_jitter_en,
    bool *sram_main_readback_locked, bool *sram_ret_readback_locked,
    bool *sram_main_readback_en, bool *sram_ret_readback_en,
    bool enable_data_ind_timing, bool *data_ind_timing_en);

/**
 * Initializes the peripherals (pinmux, uart, gpio, and timer) used by SCA code
 * and sets the sensor control module.
 *
 * @param trigger Peripheral to use for the trigger signal.
 * @param peripheral_enable Set of peripherals to enable. Remaining peripherals
 * will be disabled to reduce noise during SCA.
 * @param sensor_ctrl_enable Set the sensor control's alerts to enabled or
 * disabled.
 * @param sensor_ctrl_en_fatal Set the sensor control's alerts to fatal (true)
 * or recoverable (false).
 */
void pentest_init(
    pentest_trigger_source_t trigger, pentest_peripherals_t peripheral_enable,
    bool sensor_ctrl_enable,
    bool sensor_ctrl_en_fatal[SENSOR_CTRL_PARAM_NUM_ALERT_EVENTS]);

/**
 * Returns a handle to the intialized UART device.
 *
 * @return Handle to the initialized UART device.
 */
const dif_uart_t *pentest_get_uart(void);

/**
 * Select the capture trigger type.
 *
 * @param trigger_type The trigger type to select.
 */
void pentest_select_trigger_type(pentest_trigger_type_t trigger_type);

/**
 * Sets capture trigger high.
 *
 * The actual trigger signal used for capture is a combination (logical AND) of:
 * - the trigger gate enabled here, and
 * - the busy signal of the peripheral of interest.
 */
void pentest_set_trigger_high(void);

/**
 * Sets capture trigger low.
 */
void pentest_set_trigger_low(void);

/**
 * Functions called by `pentest_call_and_sleep()` must conform to this
 * prototype.
 */
typedef void (*sca_callee)(void);

/**
 * Calls the given function and puts Ibex to sleep.
 *
 * This function can be used to minimize noise while capturing power traces for
 * side-channel attacks, in which case `callee` would trigger the operation of
 * interest, typically a crypto operation, on the target device.
 *
 * @param callee Function to call before putting Ibex to sleep.
 * @param sleep_cycles Number of cycles to sleep.
 * @param sw_trigger Raise trigger before calling the target function.
 * @param otbn Wait until OTBN execution has finished. Not supported in
 * ENGLISH_BREAKFAST.
 */
void pentest_call_and_sleep(sca_callee callee, uint32_t sleep_cycles,
                            bool sw_trigger, bool otbn);

/**
 * Seeds the software LFSR usable e.g. for key masking.
 *
 * This functions seeds the Galois XOR type LFSR with the provided seed value.
 * Note that the LFSR itself has very low diffusion.
 *
 * @param seed The new seed value.
 * @param context Specifies which LFSR to seed.
 */
void pentest_seed_lfsr(uint32_t seed, pentest_lfsr_context_t context);

/**
 * Steps the software LFSR usable e.g. for key masking `num_steps` times.
 *
 * This function steps the Galois XOR type LFSR n times. Note that the LFSR
 * itself has very low diffusion. To improve the statistical properties, the
 * multiple steps can be run. For example, by running 32 steps it can be
 * ensured that each state bit depends on at least two previous state bits.
 *
 * @param num_steps The number of steps to run.
 * @param context Specifies which LFSR to run.

 * @return The current state of the LFSR.
 */
uint32_t pentest_next_lfsr(uint16_t num_steps, pentest_lfsr_context_t context);

/**
 * Applies a linear layer.
 *
 * This function feeds the input through a linear permutation layer. This is
 * suitable to ensure 1) that adjacent output bits of the software LFSR do not
 * go into the same S-Box (see `pentest_non_linear_layer()`) and 2) that the
 * output of S-Box(n) is not always going to be equal to the output of
 * S-Box(n+1) in the subsequent cycle. For details on how this can be achieved,
 * refer to the corresponding hardware implementation in
 * hw/ip/prim/rtl/prim_lfsr.sv.
 *
 * @param input The input value.
 *
 * @return The output of the linear layer.
 *
 */
uint32_t pentest_linear_layer(uint32_t input);

/**
 * Applies a non-linear layer.
 *
 * This function feeds the input through a non-linear layer. It is suitable to
 * improve the statistical properties of the software LFSR usable for key
 * masking, see `pentest_seed_lfsr()` and `pentest_next_lfsr()`. Internally, a
 * LUT-based AES S-Box is applied to the invididual bytes of the input word.
 *
 * In addition, the output bytes are XORed with the sbox[0]. This is useful to
 * ensure an all-zero seed (used to switch off key masking) also results in an
 * all-zero output of the non-linear layer.
 *
 * @param input The input value.
 *
 * @return The output of the non-linear layer.
 *
 */
uint32_t pentest_non_linear_layer(uint32_t input);

/**
 * Read alert information.
 *
 * This function reads the reset reason from the reason manager as well as the
 * crash dump from the reset manager. Note that on a POR this information is NOT
 * saved within the reset manager.
 *
 * @param uj An initialized uJSON context.
 * @return OK or error.
 */
status_t pentest_read_rstmgr_alert_info(ujson_t *uj);

/**
 * Return SKU config back to host.
 *
 * This function reads out different config fields of the SKU and returns them
 * back to the host over uJSON.
 *
 * @param uj An initialized uJSON context.
 * @return OK or error.
 */
status_t pentest_send_sku_config(ujson_t *uj);

/**
 * Read the Ibex RND register.
 *
 * Blocks until data is available in the Ibex RND register. Rnd data comes from
 * the EDN.
 */
uint32_t pentest_ibex_rnd32_read(void);

#endif  // OPENTITAN_SW_DEVICE_TESTS_PENETRATIONTESTS_FIRMWARE_LIB_PENTEST_LIB_H_
